import os

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langgraph.prebuilt import chat_agent_executor

# pip install dashscope
os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_PROJECT'] = 'LLMDEMO'
os.environ['LANGCHAIN_API_KEY'] = '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'
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
chatLLM = ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo",  # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    # other params...
)
# sqlalchemy 初始化MySQL数据库的连接
HOSTNAME = '192.168.1.20'
PORT = '3307'
DATABASE = 'test'
USERNAME = 'root'
PASSWORD = '123456'
# mysqlclient驱动URL
MYSQL_URI = 'mysql+pymysql://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOSTNAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)
# print(db.run("select * from student;"))
ChatPromptTemplate.from_template(
    """给定以下用户问题、SQL语句、SQL执行后的结果、回答用户问题
    Question:{question}
    SQL Query:{query}
    SQL Result:{result}
    回答:"""
)
# executor_sql_tool = QuerySQLDataBaseTool(db=db)
# test_chain = create_sql_query_chain(chatLLM, db)
# 这种只是返回了sql语句
# print(test_chain.invoke({"question": "有多少学生？"}))

toolkit = SQLDatabaseToolkit(llm = chatLLM,db = db)
tools = toolkit.get_tools()
# 使用agent完整整个数据库的整合
system_prompt = """
您是一个被设计用来与SQL数据库交互的代理。
给定一个输入问题，创建一个语法正确的SQL语句并执行，然后查看查询结果并返回答案。
除非用户指定了他们想要获得的示例的具体数量，否则始终将SQL查询限制为最多10个结果。
你可以按相关列对结果进行排序，以返回MySQL数据库中最匹配的数据。
您可以使用与数据库交互的工具。在执行查询之前，你必须仔细检查。如果在执行查询时出现错误，请重写查询SQL并重试。
不要对数据库做任何DML语句(插入，更新，删除，删除等)。

首先，你应该查看数据库中的表，看看可以查询什么。
不要跳过这一步。
然后查询最相关的表的模式。
"""
system_message = SystemMessage(content=system_prompt)
agent_executor = chat_agent_executor.create_tool_calling_executor(model=chatLLM,tools=tools,prompt=system_message)
resp = agent_executor.invoke({"messages": [HumanMessage(content="所有同学的学生编号、学生姓名、选课总数、所有课程的总成绩")]})
result = resp['messages']
print(result)
print(result[len(result)-1].content)